
"""
read and encode GenBank taxonomy records
"""

DATA_DIRECTORY = "/var/data/taxonomy"
ARCHIVE_DIRECTORY = "/var/data/taxonomy/archive"

import os
from whiff import resolver

DEBUG = False
#DEBUG = True

class DictGetter:
    def __init__(self, directory):
        self.directory = directory
    def getDicts(self, filename, columns, splitter="|"):
        print "getting filename=", filename
        filepath = os.path.join(self.directory, filename)
        lines = file(filepath).readlines()
        L = []
        count = 0
        for line in lines:
            count +=1
            if DEBUG and count>1000: break
            line = line.strip()
            print line
            if line:
                D = {}
                colvalues = line.split(splitter)
                pairs = zip(columns, colvalues)
                for (name, value) in pairs:
                    D[name] = value.strip()
                if D:
                    L.append(D)
        return L

def taxonomyDictionary(directory):
    DG = DictGetter(directory)
    result = {}
    for nameInfo in DG.getDicts("names.dmp", ["tax_id", "name_txt", "unique_name", "name_class"]):
        # tag the name with its class
        if nameInfo.has_key("name_class"):
            if nameInfo.has_key("name_txt"):
                nameInfo["name_txt"] += "["+nameInfo["name_class"]+"]"
            del nameInfo["name_class"]
        tax_id = nameInfo["tax_id"]
        if result.has_key(tax_id):
            tax_info = result[tax_id]
            tax_info.addInfo(tax_id, info=nameInfo)
        else:
            tax_info = TaxonomyInfo(tax_id, info=nameInfo)
        result[tax_id] = tax_info
    for nodeInfo in DG.getDicts("nodes.dmp", ["tax_id", "parent_id", "rank", "embl", "division"]):
        tax_id = nodeInfo["tax_id"]
        parent_id = nodeInfo["parent_id"]
        if DEBUG and not result.has_key(tax_id):
            continue
        tax_info = result[tax_id]
        tax_info.addInfo(tax_id, parent_id=parent_id, info=nodeInfo)
        parent_info = result.get(parent_id)
        if parent_info is not None:
            parent_info.addInfo(parent_id, child=tax_id)
        else:
            print "NO PARENT!!", tax_id, parent_id
    return result

class TaxonomyInfo:
    def __init__(self, tax_id, parent_id=None, info=None):
        self.tax_id = self.parent_id = None
        self.children = []
        self.info = {}
        self.addInfo(tax_id, parent_id=parent_id, info=info)
    def dictionary(self):
        D = self.info.copy()
        if D.has_key("tax_id"):
            del D["tax_id"]
        D["children"] = self.children
        return D
    def __repr__(self):
        L = []
        L.append("TaxonomyInfo("+repr(self.tax_id)+",")
        L.append("    parent="+repr(self.parent_id)+",")
        L.append("    children=[")
        for c in self.children:
            L.append("        "+repr(c)+",")
        L.append("]")
        L.append("    info="+repr(self.info))
        L.append(")")
        return "\n".join(L)
    def addInfo(self, tax_id, parent_id=None, info=None, child=None):
        assert self.tax_id is None or self.tax_id==tax_id, "tax id's don't match "+repr((self.tax_id, tax_id))
        assert self.parent_id is None or parent_id is None or self.parent_id==parent_id, "parents don't match "+repr((tax_id, parent_id, self.parent_id))
        self.tax_id = tax_id
        if parent_id is not None:
            self.parent_id = parent_id
        if child is not None and child not in self.children:
            self.children.append(child)
        if info:
            i = self.info
            for (n,v) in info.items():
                L = i.get(n) or []
                if v and v not in L:
                    L.append(v)
                i[n] = L

class TreeResource:
    def __init__(self, archive=ARCHIVE_DIRECTORY, session=None):
        self.archive = archive
        self.session = session
    def localize(self, env):
        from nucular import Nucular
        session = Nucular.Nucular(self.archive)
        return self.__class__(self.archive, session)
    def get(self, pathlist):
        assert len(pathlist)==0, "no path info expected"
        return TreeDict(self.session)

class SearchResource(TreeResource):
    def get(self, pathlist):
        assert len(pathlist)==0, "no path info expected"
        return self.session

class TreeDict:
    def __init__(self, session):
        self.session = session
    def __getitem__(self, genBankId):
        session = self.session
        if not session.hasIdentity(genBankId):
            raise KeyError, "no such gen bank entry "+repr(genBankId)
        entry = self.session.describe(genBankId)
        #pr "for", genBankId, "got entry", entry
        result = {}
        result.update(entry.attrDict())
        parent = None
        if result.has_key("parent_id"):
            parent = result["parent_id"][0]
        if parent and genBankId!=parent:
            result["parent"] = parent
        if result.has_key("unique_name"):
            names = result["unique_name"]
            #name = "<br> | ".join(result["unique_name"])
        else:
            names = result.get("name_txt", ["???"])
            #name = " / <br>".join(result.get("name_txt", ["???"]))
        if len(names)>3:
            names = names[:3]+["..."]
        names = [resolver.quote(n) for n in names]
        name = " <br> &nbsp; | ".join(names)
        if result.has_key("rank"):
            rank = result["rank"][0]
            if rank!="no rank":
                name += ": <em>"+rank+"</em>"
        #name = resolver.quote(name)
        result["body"] = formatBody(result, genBankId, name, session)
        # hack to fix root node
        if result.has_key("children") and genBankId in result["children"]:
            result["children"].remove(genBankId)
        return result

def formatBody(result, gid, name, session):
    divid = '"gbdiv%s"' % (gid)
    L = []
    #showaction = 'document.getElementById(%s).style.visibility="visible"' % divid
    showaction = "showdiv(%s)" % divid
    #hideaction = 'document.getElementById(%s).style.visibility="hidden"' % divid
    hideaction = "hidediv(%s)" % divid
    L.append ("<a href='javascript:%s'>%s</a>" % (showaction, name))
    L.append ('<div id=%s style="display:none">' % divid)
    L.append ("<table border>")
    L.append ("<tr><th colspan='2' align='right'><a href='javascript:%s'>X</a><br><th></tr>" % hideaction)
    its = result.items()
    its.sort()
    for (n,v) in its:
        if n=="parent":
            continue
        L.append("<tr> <th> %s</th>" % n)
        L.append("<td>")
        for x in v:
            if n=="children":
                x = name_of(x, session)
            else:
                x = resolver.quote(x)
            L.append( x+ "<br>")
        L.append("</td></tr>")
    L.append ("</table>")
    L.append ('</div>')
    return "\n".join(L)

def name_of(gid, session):
    entry = session.describe(gid)
    dict = entry.asDictionary()
    name = dict.get("unique_name") or dict.get("name_txt")
    name = resolver.quote(str(name))
    anchor = '<a href="javascript:focusTree(\'%s\')">%s</a>' % (gid, name)
    return anchor

def archive(fromDirectory=DATA_DIRECTORY, toNewArchive=ARCHIVE_DIRECTORY):
    from nucular import Nucular
    from nucular import entry
    session = Nucular.Nucular(toNewArchive)
    session.create()
    D = taxonomyDictionary(fromDirectory)
    its = D.items()
    its.sort()
    count = 0
    print "now storing taxonomy"
    for (n,v) in its:
        count+=1
        d = v.dictionary()
        e = entry.Entry(n,d)
        session.index(e)
        if count%10000==0:
            print "storing at", (count, n)
            session.store(lazy=True)
            session = Nucular.Nucular(toNewArchive)
    session.store(lazy=True)
    print "aggregating"
    session.aggregateRecent(verbose=True)
    print "hardening"
    session.moveTransientToBase(verbose=True)
    session.cleanUp()

def test():
    D = taxonomyDictionary(DATA_DIRECTORY)
    items = D.items()
    items.sort()
    for (n, v) in items:
        print n
        print v

if __name__=="__main__":
    #test()
    archive()

    
